import json
import time
import traceback
import modules.scripts as scripts
from modules.paths_internal import models_path, data_path
from modules import script_callbacks, shared
from PIL import Image
import numpy as np
import gradio as gr
import requests
import os
import re
import subprocess
from bs4 import BeautifulSoup
import threading

API_URL = "https://api.tzone03.xyz/"
ONLINE_DOCS_URL = API_URL + "docs/"
RESULT_PATH = "tmp/model-downloader-cn.log"
VERSION = "2.0.0"
headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/117.0.0.0 Safari/537.36"
}
libolibo_headers = {
    "token": "e9abc9c4db6e41eba7642714bdfc4c1f",
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/117.0.0.0 Safari/537.36",
    "webid": "1699929075245kztgyghw",
}
proxies = {
    'http': 'http://127.0.0.1:8888',
    'https': 'http://127.0.0.1:8888'
}


def check_aria2c():
    try:
        subprocess.run("aria2c", stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        return True
    except FileNotFoundError:
        return False


def process_image(url):
    response = requests.get(url, stream=True)
    image = Image.open(response.raw)
    return image


def get_model_path(model_type, download_type):
    if download_type == "autodl共享盘":
        co = shared.cmd_opts
        pj = os.path.join
        MODEL_TYPE_DIR = {
            "Checkpoint": ["ckpt_dir", pj("/root/autodl-fs/Stable-diffusion/models", 'Stable-diffusion')],
            "LORA": ["lora_dir", pj("/root/autodl-fs/Stable-diffusion/models", 'Lora')],
            "TextualInversion": ["embeddings_dir", pj("/root/autodl-fs/Stable-diffusion", 'embeddings')],
            "Hypernetwork": ["hypernetwork_dir", pj("/root/autodl-fs/Stable-diffusion/models", 'hypernetworks')],
            # "AestheticGradient": "",
            # "Controlnet": "", #controlnet-dir
            "LoCon": ["lyco_dir", pj("/root/autodl-fs/Stable-diffusion/models", 'LyCORIS')],
            "VAE": ["vae_dir", pj("/root/autodl-fs/Stable-diffusion/models", 'VAE')],
        }

        dir_list = MODEL_TYPE_DIR.get(model_type)
        if dir_list == None:
            return None

        if hasattr(co, dir_list[0]) and getattr(co, dir_list[0]):
            return getattr(co, dir_list[0])
        else:
            return dir_list[1]
    else:
        co = shared.cmd_opts
        pj = os.path.join
        MODEL_TYPE_DIR = {
            "Checkpoint": ["ckpt_dir", pj(models_path, 'Stable-diffusion')],
            "LORA": ["lora_dir", pj(models_path, 'Lora')],
            "TextualInversion": ["embeddings_dir", pj(data_path, 'embeddings')],
            "Hypernetwork": ["hypernetwork_dir", pj(models_path, 'hypernetworks')],
            # "AestheticGradient": "",
            # "Controlnet": "", #controlnet-dir
            "LoCon": ["lyco_dir", pj(models_path, 'LyCORIS')],
            "VAE": ["vae_dir", pj(models_path, 'VAE')],
        }

        dir_list = MODEL_TYPE_DIR.get(model_type)
        if dir_list == None:
            return None

        if hasattr(co, dir_list[0]) and getattr(co, dir_list[0]):
            return getattr(co, dir_list[0])
        else:
            return dir_list[1]


def request_civitai_detail(url):
    pattern = r'https://civitai\.com/models/(.+)'
    m = re.match(pattern, url)
    if not m:
        return False, "不是一个有效的 civitai 模型页面链接，暂不支持"

    req_url = API_URL + "civitai/models/" + m.group(1)
    res = requests.get(req_url)
    if res.ok:
        ok = True
        resp = res.json()
    else:
        ok = False
        resp = res.text
    if res.status_code >= 500:
        ok = False
        resp = "呃 服务好像挂了，理论上我应该在修了，可以进群看看进度……"
    if res.status_code >= 400:
        ok = False
        resp = "不是一个有效的 civitai 模型页面链接，暂不支持"
    if not ok:
        return [resp] + resp_to_components(None) + [gr.update(interactive=False)]

    has_download_file = False
    more_guides = ""
    if resp["version"]["file"]["downloadUrl"]:
        has_download_file = True
        more_guides = f'，点击下载按钮\n{resp["version"]["file"]["name"]}'

    return [f"预览成功{more_guides}"] + resp_to_components(resp) + \
        [gr.update(interactive=has_download_file)]


def request_libolibo_detail(url, source_type, token):
    # https://liblib.ai/modelinfo/a22ee68ab3a14461a0a916fa327b55c7
    pattern = r'(.*?)modelinfo/(.+)'
    m = re.match(pattern, url)
    if not m:
        return ["不是一个有效的 哩布哩布 模型页面链接，暂不支持"] + resp_to_components_libolibo(None) + [
            gr.update(interactive=False)]
    try:
        res = requests.get(url, headers=libolibo_headers)
        if res.status_code == 200:
            html = BeautifulSoup(res.text, "html.parser")
            data = html.find(name="script", attrs={"id": "__NEXT_DATA__"})
            json_obj = json.loads(data.text)
            has_download_file = False
            more_guides = ""
            resp = json_obj
            if "props" in json_obj and json_obj["props"] and "pageProps" in json_obj["props"] and json_obj["props"][
                "pageProps"]:
                try:
                    if resp["props"]["pageProps"]["modelData"]["versions"][0]["attachment"]["modelSource"]:
                        has_download_file = True
                        more_guides = f'，点击下载按钮\n{resp["props"]["pageProps"]["modelData"]["versions"][0]["attachment"]["modelSource"]}'
                except:
                    print(traceback.format_exc())

                return [f"预览成功{more_guides}"] + resp_to_components_libolibo(resp) + \
                    [gr.update(interactive=has_download_file)]
            else:
                header = libolibo_headers
                header.update({"token": token})
                header.update({"Referer": m.group(0)})
                ts = int(time.time())
                res = requests.post(
                    f"https://liblib-api.vibrou.com/api/www/model/getByUuid/{m.group(2)}?timestamp={ts}",
                    headers=libolibo_headers)
                if res.status_code == 200:
                    json_obj = res.json()
                    if json_obj["code"] == 0 and "data" in json_obj and json_obj["data"]:
                        try:
                            if json_obj["data"]["versions"][0]["attachment"]["modelSource"]:
                                has_download_file = True
                                more_guides = f'，点击下载按钮\n{json_obj["data"]["versions"][0]["attachment"]["modelSource"]}'
                        except:
                            print(traceback.format_exc())
                        img = json_obj["data"]["versions"][0]["imageGroup"]["coverUrl"]
                        print(json_obj)
                        return [f"预览成功{more_guides}"] + [
                            json_obj["data"]["name"],
                            "Checkpoint",
                            None,
                            json_obj["data"]["name"],
                            None,
                            None,
                            json_obj["data"]["remark"],
                            img,
                            json_obj["data"]["versions"][0]["attachment"]["modelSourceName"],
                            json_obj["data"]["versions"][0]["attachment"]["modelSource"],
                        ] + \
                            [gr.update(interactive=has_download_file)]


    except:
        print(traceback.format_exc())

    return ["失败"] + resp_to_components_libolibo(None) + [gr.update(interactive=False)]


def request_tusi_detail(url):
    # https://tusiart.com/models/620069114351768199

    pattern = r'https://tusiart.com/models/(.+)'
    m = re.match(pattern, url)
    if not m:
        return ["不是一个有效的 哩布哩布 模型页面链接，暂不支持"] + resp_to_components_libolibo(None) + [
            gr.update(interactive=False)]
    try:
        print(requests.get(url, proxies=proxies, verify=False).text)
        res = requests.get(f"https://api.tusiart.cn/community-web/v1/model/detail?modelId={m.group(1)}",
                           headers=headers)
        if res.status_code == 200:
            obj = res.json()
            if obj["code"] == "0":
                return resp_to_components_tusi(obj)


    except:
        print(traceback.format_exc())

    return ["失败"] + resp_to_components_libolibo(None) + [gr.update(interactive=False)]

    # if res.status_code >= 500:
    #     return False, "呃 服务好像挂了，理论上我应该在修了，可以进群看看进度……"
    # if res.status_code >= 400:
    #     return False, "不是一个有效的 civitai 模型页面链接，暂不支持"
    # print(res.json())
    #
    # if res.ok:
    #     return True, res.json()
    # else:
    #     return False, res.text


def resp_to_components_tusi(resp):
    if resp == None:
        return [None, None, None, None, None, None, None, None, None, None]
    file_id = resp["data"]["model"]["files"][0]["id"]
    file_url = f"https://api.tusiart.cn/community-web/v1/model/file/url?modelFileId={file_id}"
    res = requests.get(file_url, headers=headers, proxies=proxies, verify=False)
    if res.status_code == 200:
        download_url = res.json()["data"]["url"]
        has_download_file = True
        more_guides = f'，点击下载按钮\n{download_url}'
        img = resp["data"]["model"]["cover"][0]["url"]
        if img:
            img = process_image(img)
        return [f"预览成功{more_guides}"] + [
            resp["data"]["project"]["name"],
            "Checkpoint",
            None,
            resp["data"]["project"]["name"],
            None,
            None,
            resp["data"]["project"]["description"],
            img,
            resp["data"]["model"]["files"][0]["name"],
            download_url,
        ] + [gr.update(interactive=has_download_file)]
    else:
        return ["预览失败", None, None, None, None, None, None, None, None, None]


def resp_to_components_libolibo(resp):
    if resp == None:
        return [None, None, None, None, None, None, None, None, None, None]

    img = resp["props"]["pageProps"]["modelData"]["versions"][0]["imageGroup"]["images"][0]["imageUrl"]
    if img:
        img = process_image(img)

    return [
        resp["props"]["pageProps"]["modelData"]["name"],
        "Checkpoint" if resp["props"]["pageProps"]["modelData"]["checkpointType"] == 1 else "LORA",
        None,
        resp["props"]["pageProps"]["author"]["userDetail"]["nickname"],
        None,
        None,
        resp["props"]["pageProps"]["modelData"]["remark"],
        img,
        resp["props"]["pageProps"]["modelData"]["versions"][0]["attachment"]["modelSourceName"],
        resp["props"]["pageProps"]["modelData"]["versions"][0]["attachment"]["modelSource"],
    ]


def resp_to_components(resp):
    if resp == None:
        return [None, None, None, None, None, None, None, None, None, None]

    img = resp["version"]["image"]["url"]
    if img:
        img = process_image(img)

    return [
        resp["name"],
        resp["type"],
        ", ".join(resp["version"]["trainedWords"]),
        resp["creator"]["username"],
        ", ".join(resp["tags"]),
        resp["version"]["updatedAt"],
        resp["description"],
        img,
        resp["version"]["file"]["name"],
        resp["version"]["file"]["downloadUrl"],
    ]


def preview(url, source, source_type, token):
    if source == "哩布哩布":
        return request_libolibo_detail(url, source_type, token)
    elif source == "土司":
        return request_tusi_detail(url)
    else:
        return request_civitai_detail(url)


def download(model_type, filename, url, image_arr, download_type, source_type):
    if not (model_type and url and filename):
        return "下载信息缺失"
    if not source_type:
        target_path = get_model_path(model_type, download_type)
    else:
        target_path = get_model_path(source_type, download_type)
    # target_path = "D:\workspace\sd-downloader\models"
    if not target_path and source_type:
        return f"暂不支持这种类型：{source_type}"
    elif not target_path:
        return f"暂不支持这种类型：{model_type}"
    if isinstance(image_arr, np.ndarray) and image_arr.any() is not None:
        image_filename = filename.rsplit(".", 1)[0] + ".jpeg"
        target_file = os.path.join(target_path, image_filename)
        if not os.path.exists(target_file):
            image = Image.fromarray(image_arr)
            image.save(target_file)

    target_file = os.path.join(target_path, filename)
    if os.path.exists(target_file):
        return f"已经存在了，不重复下载：\n{target_file}"

    cmd = f'curl -o "{target_file}" "{url}" 2>&1'
    if check_aria2c():
        cmd = f'aria2c -c -x 16 -s 16 -k 1M -d "{target_path}" -o "{filename}" "{url}" 2>&1'

    result = subprocess.run(
        cmd,
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        encoding="UTF-8"
    )
    status_output = ""
    if result.returncode == 0:
        status_output = f"下载成功，保存到：\n{target_file}\n{result.stdout}"
    else:
        status_output = f"下载失败了，错误信息：\n{result.stdout}"

    return status_output


def request_docs():
    banner = "## https://gitee.com/hwzlikewyh/sd-downloader.git"
    footer = "## 模型下载器"

    return banner, footer


def create_czone_interface():
    czone_window = gr.Interface(fn=lambda: "这是C站界面", live=True)
    return czone_window


def create_bilibili_interface():
    bilibili_window = gr.Interface(fn=lambda: "这是哩布哩布界面", live=True)
    return bilibili_window


def create_tos_interface():
    tos_window = gr.Interface(fn=lambda: "这是土司界面", live=True)
    return tos_window


Chioes = ["Checkpoint", "LORA", "TextualInversion", "Hypernetwork", "LoCon", "VAE"]
DownloadChioes = ["SD本地", "autodl共享盘"]


def on_ui_tabs():
    banner, footer = request_docs()

    with gr.Blocks() as ui_component:
        gr.Markdown(banner)
        with gr.Row() as input_component:
            with gr.Column():
                inp_url = gr.Textbox(
                    label="模型的页面地址，不是下载链接",
                    placeholder="类似 https://civitai.com/models/28687/pen-sketch-style"
                )
                with gr.Row():
                    source = gr.Radio(["c站", "哩布哩布", "土司"], label="来源网址", info="请选择来源网址")
                    source.value = "c站"
                with gr.Row():
                    source_type = gr.Radio(Chioes, label="下载类型", info="请选择类型", visible=False, interactive=True)
                    download_type = gr.Radio(DownloadChioes, label="下载地址", info="请选择类型", visible=True,
                                             interactive=True, value=DownloadChioes[0])
                with gr.Row():
                    token = gr.Textbox("e9abc9c4db6e41eba7642714bdfc4c1f", label="哩布哩布的token",
                                       placeholder="类似 https://civitai.com/models/28687/pen-sketch-style",visible=True,interactive=True)
                with gr.Row():
                    preview_btn = gr.Button("预览")
                    download_btn = gr.Button("下载", interactive=False)
                with gr.Row():
                    result = gr.Textbox(
                        # value=result_update,
                        label="执行结果",
                        interactive=False,
                        # every=1,
                    )
            with gr.Column() as preview_component:
                with gr.Row():
                    with gr.Column() as model_info_component:
                        name = gr.Textbox(label="名称", interactive=False)
                        model_type = gr.Textbox(label="类型", interactive=False)
                        trained_words = gr.Textbox(label="触发词", interactive=False)
                        creator = gr.Textbox(label="作者", interactive=False)
                        tags = gr.Textbox(label="标签", interactive=False)
                        updated_at = gr.Textbox(label="最近更新时间", interactive=False)
                    with gr.Column() as model_image_component:
                        image = gr.Image(
                            show_label=False,
                            interactive=False,
                        )
                with gr.Accordion("介绍", open=False):
                    description = gr.HTML()
        with gr.Row(visible=False):
            filename = gr.Textbox(
                visible=False,
                label="model_filename",
                interactive=False,
            )
            download_url = gr.Textbox(
                visible=False,
                label="model_download_url",
                interactive=False,
            )

        def on_radio_change(choice):
            if choice == "c站":
                return {source_type: gr.update(visible=False, interactive=True), }  # 隐藏正常输出，显示错误框
            elif choice == "哩布哩布":
                return {
                    source_type: gr.update(value="Checkpoint", visible=True, interactive=True)
                }  # 隐藏正常输出，显示错误框
            else:
                return {source_type: gr.update(value="Checkpoint", visible=True, interactive=True), }  # 隐藏正常输出，显示错误框

        with gr.Row():
            gr.Markdown(f"版本：{VERSION}\n\n作者：@hwzlikewyh\n{footer}")

        def preview_components():
            return [
                name,
                model_type,
                trained_words,
                creator,
                tags,
                updated_at,
                description,
                image,
            ]

        def file_info_components():
            return [
                filename,
                download_url,
            ]

        preview_btn.click(
            fn=preview,
            inputs=[inp_url, source, source_type, token],
            outputs=[result] + preview_components() + \
                    file_info_components() + [download_btn]
        )
        source.change(fn=on_radio_change,
                      inputs=[source], outputs=[source_type])
        download_btn.click(
            fn=download,
            inputs=[model_type] + file_info_components() + [image] + [download_type] + [source_type],
            outputs=[result]
        )
    return [(ui_component, "模型下载", "model_downloader_cn_tab")]


# on_ui_tabs()
script_callbacks.on_ui_tabs(on_ui_tabs)
